#include <iostream>

#include <vector>
#include <map>
#include <set>
#include <memory>
#include <queue>

#include <cassert>

#include <sds/sfile.h>
#include <sds/sdataset.h>
#include <sds/sdataview.h>

#include <cluster.h>
#include <segment.h>
#include <matrix.h>

#include <lemon/smart_graph.h>
#include <lemon/grid_graph.h>
#include <lemon/matching.h>
#include <lemon/random.h>
#include <lemon/graph_to_eps.h>
#include <lemon/lgf_writer.h>
#include <lemon/kruskal.h>

const double LIMIT = 0.95;

using namespace lemon;
using namespace std;
using namespace sds;

typedef SmartGraph Graph;
GRAPH_TYPEDEFS(Graph);

typedef SDataView<int> SegmentView;
typedef SDataView<Point> PointView;
typedef SDataView<unsigned char> CharView;
typedef SDataView<std::pair<Node, Node> > NodeView;
typedef SDataView<Edge> EdgeView;

struct Core {
  GridGraph::Node node;
  double min;
  int size;
  Core(GridGraph::Node n, double m, int s)
    : node(n), min(m), size(s) {}
};

bool operator<(const Core& left, const Core& right) {
  return left.size < right.size;
}

void Read(const string& filename, GridGraph& graph,
	  GridGraph::EdgeMap<double>& weight, vector<SInteger>& size) {
  CharView source[CH];
  SFile file(filename);
  for (int k = 0; k < CH; ++k) {
    source[k] = file.getDataset(k).getDataView<unsigned char>();
  }
  
  size = source[0].getSize();
 
  Cluster dist;
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      Point point;
      for (int k = 0; k < CH; ++k) {
	point[k] = source[k][i][j];
      }
      dist += point;
    }
  }

  Matrix cm(CH, CH);
  for (int i = 0; i < CH; ++i) {
    for (int j = 0; j < CH; ++j) {
      cm(i, j) = dist.covariance(i, j);
    }
    cm(i, i) += 0.01;
  }
  Matrix cmi = cm.inverse();

  graph.resize(size[0], size[1]);

  for (GridGraph::EdgeIt it(graph); it != INVALID; ++it) {
    Vector dv(CH);
    for (int k = 0; k < CH; ++k) {
      dv(k) = (double)source[k][graph.col(graph.u(it))][graph.row(graph.u(it))]-
	(double)source[k][graph.col(graph.v(it))][graph.row(graph.v(it))];
    }
    weight[it] = exp(- (dv * cmi * dv) / CH);
  }
}

class Splitter {

  const GridGraph& _grid_graph;
  const GridGraph::EdgeMap<double>& _grid_weight;
  GridGraph::NodeMap<int>& _segments;
  GridGraph::NodeMap<char> _grid_state;

  GridGraph _dual_graph;
  GridGraph::NodeMap<std::pair<SmartGraph::Node, 
			       SmartGraph::Node> > _dual_nodes;
  GridGraph::EdgeMap<SmartGraph::Edge> _dual_edges;


  int _num;

public:

  Splitter(const GridGraph& grid_graph, 
	   const GridGraph::EdgeMap<double>& grid_weight, 
	   GridGraph::NodeMap<int>& segments) 
    : _grid_graph(grid_graph), _grid_weight(grid_weight),
      _segments(segments), _grid_state(grid_graph),
      _dual_graph(grid_graph.width() + 1, grid_graph.height() + 1),
      _dual_nodes(_dual_graph, std::pair<Node, Node>(INVALID, INVALID)),
      _dual_edges(_dual_graph, INVALID), _num(0)
      //    , _dbg_coord(_graph)
  {
    for (GridGraph::NodeIt n(_grid_graph); n != INVALID; ++n) {
      if  (_segments[n] > _num) _num = _segments[n];
    }
  }
  
  // (x,y) is always the left-top corner of a segment
  bool split(GridGraph::Node p, double lastmin, 
	     std::vector<Core>& nps) {
    if (singular(p)) return false;

    double min = numeric_limits<double>::max();
    double max = numeric_limits<double>::min();
    {

      SmartGraph treeGraph;
      SmartGraph::EdgeMap<double> treeCost(treeGraph);

      createInitGraph(p, treeGraph, treeCost);

      for (EdgeIt e(treeGraph); e != INVALID; ++e) {
	if (treeCost[e] < min) min = treeCost[e];
	if (treeCost[e] > max) max = treeCost[e];
      }
      
      SmartGraph::EdgeMap<bool> tree(treeGraph);
      kruskal(treeGraph, treeCost, tree);

      double minExternal = std::numeric_limits<double>::max();
      double maxInternal = std::numeric_limits<double>::min();
      for (EdgeIt e(treeGraph); e != INVALID; ++e) {
	if (tree[e]) {
	  if (treeCost[e] > maxInternal) {
	    maxInternal = treeCost[e];
	  }
	} else {
	  if (treeCost[e] < minExternal) {
	    minExternal = treeCost[e];
	  }	
	}
      }

      double treeMax = std::max(maxInternal, minExternal);

      if (max > maxInternal) {
	max = maxInternal;
      }

      if (min < lastmin) {
	min = lastmin;
      }
      
      clearGraph(p);
      
    }

    {

      SmartGraph matchingGraph;
      SmartGraph::EdgeMap<double> matchingWeight(matchingGraph);
      
      createGraph(p, matchingGraph, matchingWeight);

      Graph::EdgeMap<double> currentWeight(matchingGraph);
      
      while (true) {
	double cur;
	if (min > 1e-120) {
	  cur = sqrt(max * min);
	} else {
	  cur = (min + max) / 2;
	}
	std::cerr << min << ' ' << cur << ' ' << max << std::endl;
	Graph::EdgeMap<double> currentWeight(matchingGraph);
	for (EdgeIt e(matchingGraph); e != INVALID; ++e) {
	  if (matchingWeight[e] == std::numeric_limits<double>::infinity()) {
	    currentWeight[e] = 0;
	  } else { 
	    currentWeight[e] = cur - matchingWeight[e];
	  }
	}
	MaxWeightedPerfectMatching<Graph, Graph::EdgeMap<double> >
	  mwm(matchingGraph, currentWeight);
	mwm.run();
		
	double weight = mwm.matchingWeight();
	std::cerr << weight << std::endl;
	if (weight <= 0) {
	  min = cur;
	} else {
	  max = cur;
	  if (weight < 2e-5) {
	    if (cur < LIMIT) {
	      extractRegions(p, cur, mwm, nps);
	    }
	    break;
	  }
	}
	if (min >= max - 1e-20) break;
      }
      clearGraph(p);
    }
    
    return nps.size() > 0;
  }

private:

  bool singular(GridGraph::Node p) {
    for (GridGraph::OutArcIt a(_grid_graph, p); a != INVALID; ++a) {
      if (_segments[_grid_graph.target(a)] == _segments[p]) {
	return false;
      }
    }
    return true;
  }

  void createInitGraph(GridGraph::Node p, SmartGraph& treeGraph,
		       SmartGraph::EdgeMap<double>& treeCost) {
    int dn = 4;
    static const int dh[] = { 0,  1,  0, -1 };
    static const int dv[] = {-1,  0,  1,  0 };

    static const int vh[] = { 0,  1,  1,  0 };
    static const int vv[] = { 0,  0,  1,  1 };

    int seg = _segments[p];

    std::vector<GridGraph::Node> stack;
    stack.push_back(p);
    _grid_state[p] = 1;

    while (!stack.empty()) {
      GridGraph::Node pp = stack.back();
      stack.pop_back();

      for (int k = 0; k < dn; ++k) {
	GridGraph::Node np = _dual_graph(_grid_graph.col(pp) + vh[k], 
					 _grid_graph.row(pp) + vv[k]); 


	if (_dual_nodes[np].first == INVALID) {

	  Node n1 = treeGraph.addNode();
	  _dual_nodes[np].first = n1;
	}
      }

      for (int k = 0; k < dn; ++k) {

	GridGraph::Node np1 = _dual_graph(_grid_graph.col(pp) + vh[k], 
					  _grid_graph.row(pp) + vv[k]); 
	GridGraph::Node np2 = _dual_graph(_grid_graph.col(pp) + vh[(k+1) % dn], 
					 _grid_graph.row(pp) + vv[(k+1) % dn]); 
	GridGraph::Edge ne = _dual_graph.findEdge(np1, np2);

	if (_dual_edges[ne] != INVALID) continue;
	
	double dist;
	if (_grid_graph.col(pp) + dh[k] >= 0 && 
	    _grid_graph.col(pp) + dh[k] < _grid_graph.width() &&
	    _grid_graph.row(pp) + dv[k] >= 0 && 
	    _grid_graph.row(pp) + dv[k] < _grid_graph.height()) {
	  GridGraph::Node np = _grid_graph(_grid_graph.col(pp) + dh[k],
					   _grid_graph.row(pp) + dv[k]);
	  if (_segments[np] == seg) {
	    dist = _grid_weight[_grid_graph.findEdge(pp, np)];
	  } else {
	    continue;
	  }
	} else {
	  continue;
	}

	
	Node un = _dual_nodes[np1].first;
	Node vn = _dual_nodes[np2].first;
	
	Edge e = treeGraph.addEdge(un, vn);
	treeCost[e] = dist;

	_dual_edges[ne] = e;
      }


      for (GridGraph::OutArcIt a(_grid_graph, pp); a != INVALID; ++a) {
	GridGraph::Node np = _grid_graph.target(a);
	if (_segments[np] == seg && _grid_state[np] == 0) {
	  _grid_state[np] = 1;
	  stack.push_back(np);
	}
      }
    }
  }

  void createGraph(GridGraph::Node p, SmartGraph& matchingGraph,
		   SmartGraph::EdgeMap<double>& matchingWeight) {
    int dn = 4;
    static const int dh[] = { 0,  1,  0, -1 };
    static const int dv[] = {-1,  0,  1,  0 };

    static const int vh[] = { 0,  1,  1,  0 };
    static const int vv[] = { 0,  0,  1,  1 };

    int seg = _segments[p];

    std::vector<GridGraph::Node> stack;
    stack.push_back(p);
    _grid_state[p] = 1;

    while (!stack.empty()) {
      GridGraph::Node pp = stack.back();
      stack.pop_back();

      for (int k = 0; k < dn; ++k) {
	GridGraph::Node np = _dual_graph(_grid_graph.col(pp) + vh[k], 
					 _grid_graph.row(pp) + vv[k]); 


	if (_dual_nodes[np].first == INVALID) {

	  Node n1 = matchingGraph.addNode();
	  // _dbg_coord[n1] = dim2::Point<double>(np.j + 0.1, np.i + 0.1) * 15;
	  Node n2 = matchingGraph.addNode();
	  // _dbg_coord[n2] = dim2::Point<double>(np.j - 0.1, np.i - 0.1) * 15;
	  matchingWeight[matchingGraph.addEdge(n1, n2)] = 
	    std::numeric_limits<double>::infinity();
	  _dual_nodes[np].first = n1;
	  _dual_nodes[np].second = n2;
	}
      }

      for (int k = 0; k < dn; ++k) {

	GridGraph::Node np1 = _dual_graph(_grid_graph.col(pp) + vh[k], 
					  _grid_graph.row(pp) + vv[k]); 
	GridGraph::Node np2 = _dual_graph(_grid_graph.col(pp) + vh[(k+1) % dn], 
					 _grid_graph.row(pp) + vv[(k+1) % dn]); 
	GridGraph::Edge ne = _dual_graph.findEdge(np1, np2);

	if (_dual_edges[ne] != INVALID) continue;
	
	double dist;
	if (_grid_graph.col(pp) + dh[k] >= 0 && 
	    _grid_graph.col(pp) + dh[k] < _grid_graph.width() &&
	    _grid_graph.row(pp) + dv[k] >= 0 && 
	    _grid_graph.row(pp) + dv[k] < _grid_graph.height()) {
	  GridGraph::Node np = _grid_graph(_grid_graph.col(pp) + dh[k],
					   _grid_graph.row(pp) + dv[k]);
	  if (_segments[np] == seg) {
	    dist = _grid_weight[_grid_graph.findEdge(pp, np)];
	  } else {
	    dist = std::numeric_limits<double>::infinity();
	  }
	} else {
	  dist = std::numeric_limits<double>::infinity();
	}

	
	Node un1 = _dual_nodes[np1].first;
	Node un2 = _dual_nodes[np1].second;
	Node vn1 = _dual_nodes[np2].first;
	Node vn2 = _dual_nodes[np2].second;
	
	Node un = matchingGraph.addNode();
	// _dbg_coord[un] = dim2::Point<double>(0.6 * np1.j + 0.4 * np2.j, 
	// 0.6 * np1.i + 0.4 * np2.i) * 15;
	Node vn = matchingGraph.addNode();
	// _dbg_coord[vn] = dim2::Point<double>(0.4 * np1.j + 0.6 * np2.j, 
	// 0.4 * np1.i + 0.6 * np2.i) * 15;
	
	matchingWeight[matchingGraph.addEdge(un, un1)] = dist;
	matchingWeight[matchingGraph.addEdge(un, un2)] = dist;
	
	matchingWeight[matchingGraph.addEdge(vn, vn1)] = dist;
	matchingWeight[matchingGraph.addEdge(vn, vn2)] = dist;
	
	Edge e = matchingGraph.addEdge(un, vn);
	matchingWeight[e] = std::numeric_limits<double>::infinity();

	_dual_edges[ne] = e;
      }


      for (GridGraph::OutArcIt a(_grid_graph, pp); a != INVALID; ++a) {
	GridGraph::Node np = _grid_graph.target(a);
	if (_segments[np] == seg && _grid_state[np] == 0) {
	  _grid_state[np] = 1;
	  stack.push_back(np);
	}
      }
    }
  }

  void extractRegions(GridGraph::Node p, double cur,
  		      MaxWeightedPerfectMatching<Graph, 
  		      Graph::EdgeMap<double> >& mwm,
  		      std::vector<Core>& nps) {
    static const int dn = 4;
    
    static const int dv[] = {-1,  0,  1,  0 };
    static const int dh[] = { 0,  1,  0, -1 };

    static const int vv[] = { 0,  0,  1,  1 };
    static const int vh[] = { 0,  1,  1,  0 };

    static const int ed[] = { 0,  1,  1,  0 };
    static const bool eb[] = { true, false, true, false };
    
    int seg = _segments[p];

    std::vector<GridGraph::Node> sources;
    sources.push_back(p);
    _grid_state[p] = 2;
    
    while (!sources.empty()) {
      p = sources.back();
      sources.pop_back();
      if (_grid_state[p] == 3) continue;

      std::vector<GridGraph::Node> stack;
      stack.push_back(p);
      _grid_state[p] = 3;
      ++_num;

      int ss = 0;
      
      while (!stack.empty()) {
  	GridGraph::Node pp = stack.back();
  	stack.pop_back();
  	_segments[pp] = _num;
	++ss;
	
  	for (int k = 0; k < 4; ++k) {
	  if (_grid_graph.col(pp) + dh[k] >= 0 && 
	      _grid_graph.col(pp) + dh[k] < _grid_graph.width() &&
	      _grid_graph.row(pp) + dv[k] >= 0 && 
	      _grid_graph.row(pp) + dv[k] < _grid_graph.height()) {
	    GridGraph::Node np = _grid_graph(_grid_graph.col(pp) + dh[k], 
					     _grid_graph.row(pp) + dv[k]);

	    if (_segments[np] == seg) {
	      GridGraph::Node np1 = 
		_dual_graph(_grid_graph.col(pp) + vh[k], 
			    _grid_graph.row(pp) + vv[k]); 
	      GridGraph::Node np2 = 
		_dual_graph(_grid_graph.col(pp) + vh[(k+1) % dn], 
			    _grid_graph.row(pp) + vv[(k+1) % dn]); 
	      SmartGraph::Edge e = _dual_edges[_dual_graph.findEdge(np1, np2)];
	      if (!mwm.matching(e)) {
		if (_grid_state[np] == 1) {
		  sources.push_back(np);
		  _grid_state[np] = 2;
		}
	      } else {
		if (_grid_state[np] == 1 || _grid_state[np] == 2) {
		  stack.push_back(np);
		  _grid_state[np] = 3;
		}
	      }
	    }
  	  }
  	}
      }
      std::cerr << "Push back: " << ' '
		<< _grid_graph.col(p) << ' ' << _grid_graph.row(p) << ' '
		<< cur << ' ' << ss << std::endl;
      nps.push_back(Core(p, cur, ss));
    }
  }

  void clearGraph(GridGraph::Node p) {
    static const int dn = 4;
    
    static const int dv[] = {-1,  0,  1,  0 };
    static const int dh[] = { 0,  1,  0, -1 };
    
    static const int vv[] = { 0,  0,  1,  1 };
    static const int vh[] = { 0,  1,  1,  0 };

    static const int ed[] = {0,  1,  1,  0 };
    static const bool eb[] = { true, false, true, false };

    int seg = _segments[p];

    std::vector<GridGraph::Node> stack;
    stack.push_back(p);
    _grid_state[p] = 0;

    while (!stack.empty()) {
      GridGraph::Node pp = stack.back();
      stack.pop_back();

      for (int k = 0; k < dn; ++k) {
	GridGraph::Node np = _dual_graph(_grid_graph.col(pp) + vh[k], 
					 _grid_graph.row(pp) + vv[k]);
	_dual_nodes[np].first = INVALID;
	_dual_nodes[np].second = INVALID;
      }

      for (int k = 0; k < dn; ++k) {
	GridGraph::Node np1 = _dual_graph(_grid_graph.col(pp) + vh[k], 
					  _grid_graph.row(pp) + vv[k]); 
	GridGraph::Node np2 = _dual_graph(_grid_graph.col(pp) + vh[(k+1) % dn], 
					 _grid_graph.row(pp) + vv[(k+1) % dn]); 
	_dual_edges[_dual_graph.findEdge(np1, np2)] = INVALID;
      }


      for (GridGraph::OutArcIt a(_grid_graph, pp); a != INVALID; ++a) {
	GridGraph::Node np = _grid_graph.target(a);
	if (_grid_state[np] > 0) {
	  _grid_state[np] = 0;
	  stack.push_back(np);
	}
      }
    }

    // graphWriter(_dual_graph, std::cerr).
    //   edgeMap("edges", composeMap(IdMap<Graph, Edge>
    //   (matchingGraph), _dual_edges)).run();
  }
};

int main(int argc, const char* argv[]) {

  std::vector<SInteger> size;
  GridGraph graph(0, 0);
  GridGraph::EdgeMap<double> weight(graph);
  GridGraph::NodeMap<int> segments(graph);

  Read(argv[1], graph, weight, size);


  std::priority_queue<Core> queue;
  //  queue.push(Core(graph(0, 0), 0.0, size[0] * size[1]));


  // const int hs = 4, vs = 4;
  // std::set<int> classes;
  // for (int i = 0; i < graph.width(); ++i) {
  //   for (int j = 0; j < graph.height(); ++j) {
  //     int ci = i * hs / graph.width();
  //     int cj = j * vs / graph.height();
  //     segments[graph(i, j)] = ci * vs + cj;
  //     if (classes.find(ci * vs + cj) == classes.end()) {
  // 	classes.insert(ci * vs + cj);
  // 	queue.push(Core(graph(i, j), 0.0, size[0] * size[1]));
  //     }
  //   }
  // }

  if (argc > 2) {
    SFile segfile(argv[2]);
    CharView cseg = segfile.getDataset(0).getDataView<unsigned char>();
    SegmentView seg(cseg.getSize());
    ReadSegments(cseg, seg);

    
    std::set<int> classes;
    for (int i = 0; i < graph.width(); ++i) {
      for (int j = 0; j < graph.height(); ++j) {
	segments[graph(i, j)] = seg[i][j];
	if (classes.find(seg[i][j]) == classes.end()) {
	  classes.insert(seg[i][j]);
	  queue.push(Core(graph(i, j), 0.0, size[0] * size[1]));
	}
      }
    }    
  } else {
    for (int i = 0; i < graph.width(); ++i) {
      for (int j = 0; j < graph.height(); ++j) {
	segments[graph(i, j)] = 0;
      }
    }
    queue.push(Core(graph(0, 0), 0.0, size[0] * size[1]));
  }
  
  Splitter splitter(graph, weight, segments);
  
  int kstart = 1000000;
  int k = kstart;
  while (!queue.empty()) {
    std::cerr << "START: " << ++k << ' ' << queue.top().min << ' '
	      << queue.top().size << std::endl;
    std::vector<Core> nps;
    std::cerr << "P: " << graph.col(queue.top().node) << ' '
	      << graph.row(queue.top().node) << ' '
	      << queue.top().min << std::endl;
    double value = splitter.split(queue.top().node, 
				  queue.top().min, nps);
    queue.pop();
    for (int i = 0; i < static_cast<int>(nps.size()); ++i) {
      queue.push(nps[i]);
    }

    if (k < kstart + 50 || k % 15 == 0 || queue.empty()) {
      SDataView<int> seg(size);
      for (int i = 0; i < size[0]; ++i) {
	for (int j = 0; j < size[1]; ++j) {
	  seg[i][j] = segments[graph(i, j)];
	}
      }
      SDataView<unsigned char> trg(size);
      
      WriteSegments(seg, trg);
    
      std::ostringstream oss;
      oss << "meancutsegmentation_" << k << ".hdf";
    
      SFile trgfile(oss.str(), SFile::CREATE);
      trgfile.createDataset("segmented", sds::UINT8, size).
	setDataView<unsigned char>(trg);
    }
    
  }
  
  
  return 0;
}

